package io.renren.sdapi.service.impl;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson2.*;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.Gson;
import io.renren.constant.BizSystemException;
import io.renren.constant.SystemConstatnt;
import io.renren.enums.ResultCodeEnum;
import io.renren.enums.ScheduleTypeEnum;
import io.renren.membermanage.dto.ModelTrainParamDto;
import io.renren.membermanage.entity.GenerateLog;
import io.renren.membermanage.entity.Keyword;
import io.renren.membermanage.service.GenerateLogService;
import io.renren.membermanage.service.KeywordService;
import io.renren.sdapi.dto.httprequestparam.QueryHistoryDto;
import io.renren.sdapi.dto.httprequestparam.SseMessageDto;
import io.renren.sdapi.dto.httprequestparam.Text2ImageRequestParam;
import io.renren.sdapi.dto.httprequestparam.image2image.ApiTypeEnum;
import io.renren.sdapi.dto.httprequestparam.image2image.BaseParam;
import io.renren.sdapi.service.SseService;
import io.renren.sdapi.service.StableDiffusionApiService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.io.*;
import java.time.LocalDateTime;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Service
public class StableDiffusionApiServiceImpl implements StableDiffusionApiService {

    @Resource
    private GenerateLogService generateLogService;

    @Autowired
    private RestTemplate restTemplate;

    @Resource
    private KeywordService keywordService;

    @Resource
    private SseService service;

    @Value("${Stable-Diffusion.master.url}")
    private String masterUrl;

    @Value("${image.sava.path}")
    private String imageSavePath;

    @Value("${model.train.api}")
    private String modelTrainApi;

    @Override
    public Map<String, List<String>> text2Image(Text2ImageRequestParam param, Long userId) throws Exception {
        JSON.config(JSONWriter.Feature.LargeObject);
        String str = getLoraStr(param.getLora());
        if (null != str){
            param.setPrompt(analyzeWord(param.getPrompt()) + "," + str);
        }else {
            param.setPrompt(analyzeWord(param.getPrompt()));
        }
        //param.setPrompt(param.getPrompt() + ",extreme detail description,aHDR, UHD, 8K,masterpiece,Highly detailed");
        param.setPrompt(param.getPrompt() + ", masterpiece, High quality");
        String text2imageUrl = masterUrl + "txt2img";
        param.setSampler_name("DPM++ 2M");
        param.setScheduler("Karras");
        ResponseEntity<String> resp = restTemplate.postForEntity(text2imageUrl, param, String.class);
        return handleData(ApiTypeEnum.TEXT2IMAGE.getCode(), param, resp, JSONObject.parseObject(JSON.toJSONString(param)).toString(), userId);
    }

    @Override
    public Map<String, List<String>> image2Image(BaseParam request, Long userId) throws Exception{
        JSON.config(JSONWriter.Feature.LargeObject);
        String image2imageUrl = masterUrl + "img2img";
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        String json = JSONObject.parseObject(JSON.toJSONString(request)).toString();
        //log.info(" ---------------------- Stable Diffusion param is ---------------- {}", json);
        HttpEntity<String> entity = new HttpEntity<>(json,headers);
        ResponseEntity<String> responseEntity = restTemplate.exchange(image2imageUrl, HttpMethod.POST, entity, String.class);
        return handleData(ApiTypeEnum.IMAGE2IMAGE.getCode(), request, responseEntity, json, userId);
    }

    @Override
    public String getLoraSsOutputName(String path){
        String ssOutputName = null;
        String loraSsOutputNameUrl = masterUrl + "loras";
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        HttpEntity<String> entity = new HttpEntity<>(headers);
        ResponseEntity<String> responseEntity = restTemplate.exchange(loraSsOutputNameUrl, HttpMethod.GET, entity, String.class);
        if (!responseEntity.getStatusCode().is2xxSuccessful()) {
            throw new BizSystemException(ResultCodeEnum.REMOTE_CALL_FAILURE.getCode(), ResultCodeEnum.REMOTE_CALL_FAILURE.getMessage());
        }
        if(null == responseEntity.getBody() || responseEntity.getBody().isEmpty()){
            log.error("800,调用sd接口sdapi/vi/loras成功但返回结果为空");
            throw new BizSystemException(ResultCodeEnum.REMOTE_CALL_FAILURE.getCode(), ResultCodeEnum.REMOTE_CALL_FAILURE.getMessage());
        }
        try{
            ObjectMapper objectMapper = new ObjectMapper();
            List<Map<String, Object>> responseList = objectMapper.readValue(responseEntity.getBody(), List.class);
            for(Map<String, Object> item : responseList){
                if(ObjectUtil.isEmpty(item.get("path"))){
                    continue;
                }
                String path1 = item.get("path").toString().replace("\\", "/");
//                log.info(path1+"-------------"+path);
                if (path1.equals(path)){
                    Map<String, Object> metadata = objectMapper.readValue(objectMapper.writeValueAsString(item.get("metadata")), Map.class);
                    if (null != metadata.get("ss_output_name")) {
                        ssOutputName = metadata.get("ss_output_name").toString();
                        break;
                    }
                }
            }
            if (StrUtil.isBlank(ssOutputName)){
                log.error("801,调用sd接口sdapi/vi/loras成功但没有与当前上传的模型相匹配的路径");
                throw new BizSystemException(ResultCodeEnum.REMOTE_CALL_RESULT_ERROR.getCode(), ResultCodeEnum.REMOTE_CALL_RESULT_ERROR.getMessage());
            }
            return ssOutputName;
        }catch (Exception e){
            e.printStackTrace();
            throw new BizSystemException(ResultCodeEnum.REMOTE_CALL_FAILURE.getCode(), ResultCodeEnum.REMOTE_CALL_FAILURE.getMessage());
        }
    }

    @Override
    public boolean refresh(){
        String refreshUrl = masterUrl + "refresh-loras";
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        HttpEntity<String> entity = new HttpEntity<>("", headers);
        ResponseEntity<String> responseEntity = restTemplate.exchange(refreshUrl, HttpMethod.POST, entity, String.class);
        if (!responseEntity.getStatusCode().is2xxSuccessful()) {
            throw new BizSystemException(ResultCodeEnum.REMOTE_CALL_REFRESH_FAILURE.getCode(), ResultCodeEnum.REMOTE_CALL_REFRESH_FAILURE.getMessage());
        }
        return true;
    }

    public String modelTrain(String modelTrainId, ModelTrainParamDto dto){
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        String json = JSONObject.parseObject(JSON.toJSONString(dto)).toString();
        log.info("Model train request param : ---------------" + json);
        HttpEntity<String> entity = new HttpEntity<>(json, headers);
        ResponseEntity<String> responseEntity = restTemplate.exchange(modelTrainApi, HttpMethod.POST, entity, String.class);
        if (null == responseEntity.getBody()){
            log.error("------------Task major key " + modelTrainId + ": There was a problem calling the interface for model training!The return value is null!");
            return null;
        }
        JSONObject result = JSONObject.parseObject(responseEntity.getBody());
        log.info("--------------sdtrain return----"+ JSONUtil.toJsonStr(result));
        if ("success".equals(result.get("status"))) {
            String message = result.get("message").toString();
            int lastIndex = message.lastIndexOf(' ');
            // 截取最后一个空格之后的部分
            String s = (lastIndex != -1) ? message.substring(lastIndex + 1) : null;
            if (StrUtil.isNotBlank(s)){
                return s;
            }
            log.error("There was a problem calling the interface for model training!");
            return null;
        }else {
            log.error("Model train error info : " + result);
            return null;
        }
    }

    @Override
    public Page<GenerateLog> getHistory(QueryHistoryDto dto) {
        QueryWrapper<GenerateLog> query = new QueryWrapper<>();
        query.orderByDesc("create_time");
        query.eq("create_by",dto.getUserId().toString());
        query.select("id", "generate_param", "create_time","api_type");
        Page<GenerateLog> page = Page.of(dto.getPageNum(), dto.getPageSize());
        Page<GenerateLog> data = generateLogService.page(page,query);
        data.getRecords().forEach(log -> {
            if(log.getApiType() == 1) {
                JSONObject param = JSONObject.parseObject(log.getGenerateParam());
                log.setGenerateParam(param.getString("prompt"));
            }else if (log.getApiType() == 2){
                JSONObject param = JSONObject.parseObject(log.getGenerateParam());
                log.setGenerateParam(param.getString("prompt"));
            }
        });
        return data;
    }

    @Override
    public GenerateLog getOneHistory(Long id) {
        return generateLogService.getById(id);
    }

    private Map<String, List<String>> handleData(Integer code, Object param, ResponseEntity<String> resp, String apiRequestParam, Long userId) {
        if (!resp.getStatusCode().is2xxSuccessful()) {
            throw new BizSystemException(ResultCodeEnum.CREATE_IMAGE_ERROR.getCode(), ResultCodeEnum.CREATE_IMAGE_ERROR.getMessage());
        }
        JSONObject json = JSONObject.parseObject(resp.getBody());

        if (ObjectUtil.isNull(json)) {
            throw new BizSystemException(ResultCodeEnum.CREATE_IMAGE_ERROR.getCode(), ResultCodeEnum.CREATE_IMAGE_ERROR.getMessage());
        }
        if (StrUtil.isBlank(json.getString("images"))) {
            throw new BizSystemException(ResultCodeEnum.CREATE_IMAGE_ERROR.getCode(), ResultCodeEnum.CREATE_IMAGE_ERROR.getMessage());
        }
        JSONArray images = json.getJSONArray("images");
        List<String> paths = new ArrayList<>();
        List<String> browses = new ArrayList<>();
        String property = System.getProperty("user.dir");
        property = property.substring(0,property.length()-1);
        property = property.substring(0,property.lastIndexOf("\\"));
        for(int i = 0; i < images.size(); i++){
            byte[] bytes = Base64.decode(images.getString(i));
            String fileName = System.currentTimeMillis() + i + ".jpg";
            File file = new File(property + imageSavePath, fileName);
            String filePath = file.getAbsolutePath();
            System.out.print(filePath);
            try(OutputStream outputStream = new FileOutputStream(filePath)){
                outputStream.write(bytes);
            }catch (IOException e){
                throw new BizSystemException(ResultCodeEnum.CREATE_IMAGE_ERROR.getCode(), ResultCodeEnum.CREATE_IMAGE_ERROR.getMessage());
            }
            paths.add(filePath);
            browses.add("/test_user_1/picture/" + fileName);
        }
        GenerateLog log = new GenerateLog();
        // 创建 Gson 实例
        Gson gson = new Gson();
        // 将对象转换为 JSON 字符串
        log.setGenerateParam(gson.toJson(param));
        log.setApiType(code);
        log.setCreateBy(userId);
        log.setCreateTime(LocalDateTime.now());
        log.setSdApiRequestParam(apiRequestParam);
        log.setSdApiResponseData(json.toString());
        boolean save = generateLogService.save(log);
        if(!save){
            throw new BizSystemException(ResultCodeEnum.CREATE_IMAGE_ERROR.getCode(), ResultCodeEnum.CREATE_IMAGE_ERROR.getMessage());
        }
        Map<String, List<String>> map = new HashMap<>();
        map.put("viewPath", paths);
        map.put("browse", browses);
        if (2 == code){
            SseMessageDto sseMessageDto = new SseMessageDto();
            sseMessageDto.setUserId(userId.toString());
            sseMessageDto.setData(map);
            sseMessageDto.setTotal(1);
            service.sendMessage(sseMessageDto);
            return new HashMap<>();
        }
        return map;
    }

    public String analyzeWord(String word){
        Map<String, Keyword> collect = keywordService.list().stream().collect(Collectors.toMap(Keyword::getWordName, c -> c, (c1, c2) -> c1));
        String replace = word.trim().replace(",+", ",");
        String[] split = replace.split(",");
        StringBuilder str = null;
        int count = 0;
        for(String s : split){
            if(null != collect.get(s)){
                if(StrUtil.isNotBlank(str)){
                    str = str.append(",").append(collect.get(s).getWordEn());
                    if (Objects.equals(collect.get(s).getKeywordTypeId(), SystemConstatnt.KEYWORD_TYPE)){
                        count = count + 1;
                    }
                }else {
                    str = new StringBuilder(collect.get(s).getWordEn());
                    count = count + 1;
                }
            }else{
                if(StrUtil.isNotBlank(str)){
                    str = str.append(",").append(s);
                }else {
                    str = new StringBuilder(s);
                }
            }
        }
        String prompt = StrUtil.removeSuffix(str.toString().replace(",,", ","), ",");
        if (count > 0){
            String weight = String.format("%.2f", (double)1 / (double)count);
            prompt = prompt.replaceAll(":1>",":" + weight + ">");
        }
        log.info("处理后的提示词：{}", prompt);
        return prompt;
    }

    public String getLoraStr(List<Map<String, String>> lora){
        StringBuilder str = null;
        if (CollUtil.isNotEmpty(lora)){
            for (Map<String, String> c : lora){
                if (null == str){
                    str = new StringBuilder();
                }
                str = str.append("<lora:");
                str = str.append(c.get("name")).append(":");
                str = str.append(c.get("weight")).append(">");
                str = str.append(",");
                str = str.append(c.get("keyword"));
                str = str.append(",");
            }
            return StrUtil.removeSuffix(str.toString().replace(",,", ","), ",");
        }
        return null;
    }
}
